//
//  NSDataEncrypt.m
//  aestest
//
//  Created by yaoziyang on 14-7-15.
//  Copyright (c) 2014年 yaoziyang. All rights reserved.
//

#import "NSDataEncrypt.h"
#import <CommonCrypto/CommonHMAC.h>
#import <CommonCrypto/CommonCryptor.h>
#import <CommonCrypto/CommonDigest.h>

@implementation NSData(AESUtilitiy)

+ (NSData*)readDataFromFile:(NSURL*)url
{
    NSError* error = nil;
    NSFileHandle* fHandle = [NSFileHandle fileHandleForReadingFromURL:url error:&error];
    if (fHandle == nil || error != nil) {
        @throw [NSException exceptionWithName:@"AESUtilitiy"
                                       reason:@"can not open file for read!"
                                     userInfo:nil];
    }
    NSData* data = [fHandle readDataToEndOfFile];
    [fHandle closeFile];
    return data;
}

- (void)checAESKey:(NSData*)key
{
    NSUInteger keySize = [key length];
    if (keySize != 16 && keySize != 24 && keySize != 32) {
        @throw [NSException exceptionWithName:@"AESUtilitiy"
                                       reason:@"invalid key length, must be 16, 24, or 32!"
                                     userInfo:nil];
    }
}

- (NSData*)aesEncrypt:(NSData*)key
{
    if (self == nil || [self length] == 0) {
        return nil;
    }
    
    [self checAESKey:key];
    
    NSUInteger keySize = [key length];
    NSUInteger dataSize = [self length];
    
    int paddingSize = kCCKeySizeAES128 - (dataSize % kCCKeySizeAES128);
    dataSize += paddingSize;
    
    void* dataBuffer = malloc(dataSize);
    memcpy(dataBuffer, [self bytes], [self length]);
    
    for (NSUInteger i = [self length]; i < dataSize; i++) {
        *(char*)(dataBuffer + i) = paddingSize;   // padding
    }
    
    // setup output buffer
    size_t bufferSize = dataSize + kCCBlockSizeAES128;
    void *buffer = malloc(bufferSize);
    memset(buffer, 0, bufferSize);
    
    char iv[16] = {0};
    size_t encryptedSize = 0;
    
    // do encrypt
    CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt,
                                          kCCAlgorithmAES128,
                                          kCCOptionECBMode,
                                          [key bytes],
                                          keySize,      // kCCKeySizeAES
                                          iv,           // IV
                                          dataBuffer,
                                          dataSize,
                                          buffer,
                                          bufferSize,
                                          &encryptedSize);
    
    free(dataBuffer);
    
    if (cryptStatus == kCCSuccess) {
        NSData* ret = [NSData dataWithBytes:buffer length:encryptedSize];
        free(buffer);
        return ret;
    }
    else {
        free(buffer);
        @throw [NSException exceptionWithName:@"AESUtilitiy"
                                       reason:@"encrypt error!"
                                     userInfo:nil];
        return nil;
    }
}

- (NSData*)aesDecrypt:(NSData*)key
{
    if (self == nil || [self length] == 0) {
        return nil;
    }
    
    [self checAESKey:key];
    
    NSUInteger dataSize = [self length];
    
    // setup output buffer
    size_t bufferSize = dataSize + kCCBlockSizeAES128;
    void *buffer = malloc(bufferSize);
    memset(buffer, 0, bufferSize);
    
    char iv[16] = {0};
    size_t decryptedSize = 0;
    
    // do decrypt
    CCCryptorStatus deCryptStatus = CCCrypt(kCCDecrypt,
                                            kCCAlgorithmAES128,
                                            kCCOptionECBMode,
                                            [key bytes],
                                            [key length],   // kCCKeySizeAES
                                            iv,             // IV
                                            [self bytes],
                                            dataSize,
                                            buffer,
                                            bufferSize,
                                            &decryptedSize);
    
    if (deCryptStatus == kCCSuccess) {
        char padding = 0;
        if (decryptedSize > 0) {
            padding = *(char*)(buffer+decryptedSize-1);
        }
        NSData* ret = [NSData dataWithBytes:buffer length:decryptedSize-padding];
        free(buffer);
        return ret;
    }
    else {
        free(buffer);
        @throw [NSException exceptionWithName:@"AESUtilitiy"
                                       reason:@"decrypt error!"
                                     userInfo:nil];
        return nil;
    }
}


+ (NSData*)aesEncryptFile:(NSURL*)url key:(NSData*)key
{
    return [[self readDataFromFile:url] aesEncrypt:key];
}

+ (NSData*)aesDecryptFile:(NSURL*)url key:(NSData*)key
{
    return [[self readDataFromFile:url] aesDecrypt:key];
}

+ (NSData*)parseHexString:(NSString*)str
{
    if (!str || str.length == 0) {
        return nil;
    }
    
    NSUInteger n = str.length/2 + (str.length%2);
    NSUInteger b = str.length - (n-1)*2;   // length of last hex string.
    
    char* buf = (char*)malloc(n);
    for (NSUInteger i = 0; i < n; i++) {
        NSString* sub = [str substringWithRange:NSMakeRange(i*2, i < n-1 ? 2:b)];
        NSScanner* scanner = [NSScanner scannerWithString:sub];
        unsigned iValue = 0;
        [scanner scanHexInt:&iValue];
        *(buf+i) = (char)iValue;
    }
    
    return [[NSData alloc] initWithBytes:buf length:n];
}

@end



@implementation NSString (md5)

-(NSData *) md5Data
{
    const char *original_str = [self UTF8String];
    unsigned char result[CC_MD5_DIGEST_LENGTH];
    CC_MD5(original_str, strlen(original_str), result);
    
    return [NSData dataWithBytes:result length:CC_MD5_DIGEST_LENGTH];
}

@end
